"""
Copyright (c) 2024 [Baijiankang/DUT]
简易火车订票系统

联系信息: [jiankang_bai@163.com]
项目地址: [https://gitee.com/BJKgitee/simple-train-ticket-booking-system]

保留所有权利。
"""

import json
from collections import defaultdict
import heapq

# 加载数据
def load_data(file_path):
    try:
        with open(file_path, "r", encoding="utf-8") as f:
            data = json.load(f)
        return data
    except FileNotFoundError:
        print(f"文件 {file_path} 未找到")
        return None
    except json.JSONDecodeError:
        print(f"文件 {file_path} 不是有效的JSON格式")
        return None

# 判断是否为子序列
def is_subsequence(sub, main):
    sub_len = len(sub)
    return any(sub == main[i : i + sub_len] for i in range(len(main) - sub_len + 1))

# 判断直达
def nonstop(data, start, end):
    shortest_time = float("inf")
    path = []
    direction = None  # 新增一个变量来存储列车的行驶方向
    train = []  # 新增一个变量来存储列车的编号

    for route in data:
        if start in route["stations"] and end in route["stations"]:
            idx1 = route["stations"].index(start)
            idx2 = route["stations"].index(end)
            if idx1 < idx2:  # 如果起始站在终点站之前，列车的行驶方向是 'forward'
                if not route["forward"]:  # 如果 'forward' 为空，直接跳过这条路线
                    continue
                for train_id, tickets in route["forward"].items():
                    if all(
                        ticket > 0 for ticket in tickets[idx1 + 1 : idx2 + 1]
                    ):  # 检查所有的票数是否都大于零
                        time = sum(route["time"][idx1 + 1 : idx2 + 1])
                        if time < shortest_time:
                            shortest_time = time
                            path = route["stations"][idx1 : idx2 + 1]
                            direction = "forward"
                            train = [train_id]
                        elif time == shortest_time:
                            train.append(train_id)

            else:  # 如果起始站在终点站之后，列车的行驶方向是 'backward'
                if not route["backward"]:  # 如果 'backward' 为空，直接跳过这条路线
                    continue
                for train_id, tickets in route["backward"].items():
                    if all(
                        ticket > 0 for ticket in tickets[idx2:idx1]
                    ):  # 检查所有的票数是否都大于零
                        time = sum(route["time"][idx2 + 1 : idx1 + 1])
                        if time < shortest_time:
                            shortest_time = time
                            path = route["stations"][idx2 : idx1 + 1]
                            path.reverse()  # 反转路径
                            direction = "backward"
                            train = [train_id]
                        elif time == shortest_time:
                            train.append(train_id)

    if path == []:
        return float("inf"), [], "", []
    else:
        return (
            shortest_time,
            path,
            direction,
            train,
        )  # 返回结果中包含列车的行驶方向和列车的编号


# 线路操作类
class Railway:
    def __init__(self, data):
        self.data = data
        self.graph_distance = defaultdict(list)
        self.graph_time = defaultdict(list)
        self.graph_direct = defaultdict(set)
        self.station_line = defaultdict(
            lambda: None
        )  # 使用 defaultdict，当键不存在时返回 None

        for idx, route in enumerate(data):  # 使用 enumerate 来获取每个路线的索引
            for i in range(len(route["stations"]) - 1):
                # 如果票数之和大于零，将这两个站点视为可达
                if route["forward"]:
                    total_tickets = sum(
                        tickets[i + 1] for train_id, tickets in route["forward"].items()
                    )
                    if total_tickets > 0:
                        self.graph_time[route["stations"][i]].append(
                            (route["stations"][i + 1], route["time"][i + 1])
                        )
                if route["backward"]:
                    total_tickets = sum(
                        tickets[i - 1]
                        for train_id, tickets in route["backward"].items()
                    )
                    if total_tickets > 0:
                        self.graph_time[route["stations"][i + 1]].append(
                            (route["stations"][i], route["time"][i + 1])
                        )
                self.graph_direct[route["stations"][i]].add(route["stations"][i + 1])
                self.graph_direct[route["stations"][i + 1]].add(route["stations"][i])
                # 在读取数据时，为每个站点添加一个标记，表示它在 JSON 数据中的位置
                self.station_line[route["stations"][i]] = idx
    
    # 添加边
    def _add_edge(self, u, v, distance, time):
        self.graph_distance[u].append((v, distance))
        self.graph_time[u].append((v, time))

    # 最短路径算法
    def shortest_path(self, start, end, graph):
        queue = [(0, start, [])]
        seen = set()
        while queue:
            dist, station, path = heapq.heappop(queue)
            if station not in seen:
                seen.add(station)
                path = path + [station]
                if station == end:
                    return dist, path
                for next_station, weight in graph.get(station, []):
                    if next_station not in seen:
                        heapq.heappush(queue, (dist + weight, next_station, path))
        return float("inf"), []

    # 最短距离路径
    def shortest_distance_path(self, start, end):
        return self.shortest_path(start, end, self.graph_distance)

    # 最短时间路径
    def shortest_time_path(self, start, end):
        return self.shortest_path(start, end, self.graph_time)

    # 查找班次
    def get_schedules(self, start, end):
        for route in self.data:
            if route['start_s'] == start and route['end_s'] == end:
                schedules = []
                schedules.extend(route['forward'].keys())
                schedules.extend(route['backward'].keys())
                return schedules
        return []


    
    # 列车路线
    def get_all_routes(self):
        routes = []
        for route in self.data:
            stations = route["stations"]
            distances = route["distance"]
            times = route["time"]
            for i in range(len(stations) - 1):
                start = stations[i]
                end = stations[i + 1]
                distance = distances[i + 1]
                time = times[i + 1]
                routes.append((start, end, distance, time))
        return routes
    
    def shortest_time_path_with_nonstop_check(self, start, end):
        # 首先检查是否有直达的路径
        direct_time, direct_path, direct_direction, direct_train = nonstop(self.data, start, end)
        if direct_time != float("inf"):
            return direct_time, direct_path, [], [direct_direction], [direct_train]

        # 如果没有直达的路径，进行换乘路径的计算
        shortest_time, short_time_path = self.shortest_time_path(start, end)
        dirc = []
        train = []
        nonstop_stations = []
        i = 0
        for j in range(len(short_time_path) - 1):
            path = nonstop(self.data, short_time_path[i], short_time_path[j + 1])[1]
            if not path or not (is_subsequence(path, short_time_path)):
                nonstop_stations.append(short_time_path[j])
                new_dirc, new_train = nonstop(self.data, short_time_path[i], short_time_path[j])[2:]
                dirc.append(new_dirc)
                train.append(new_train)
                i = j
        new_dirc, new_train = nonstop(self.data, short_time_path[i], end)[2:]
        dirc.append(new_dirc)
        train.append(new_train)
        return shortest_time, short_time_path, nonstop_stations, dirc, train
    
    def min_transfer_path(self, start, end):
        queue = [(start, [start])]
        seen = set()
        min_transfers = float("inf")
        min_transfer_path = []
        while queue:
            (station, path) = queue.pop(0)
            if station == end:
                transfers = len(
                    [
                        station
                        for i, station in enumerate(path[:-1])
                        if path[i + 1] not in self.graph_direct[station]
                    ]
                )
                if transfers < min_transfers:
                    min_transfers = transfers
                    min_transfer_path = path
            else:
                for next_station, _ in self.graph_time.get(station, []):
                    if next_station not in seen:
                        seen.add(next_station)
                        queue.append((next_station, path + [next_station]))
        return min_transfer_path

    # 最少换乘
    def min_transfer_path_with_nonstop_check(self, start, end):
        min_transfer_path = self.min_transfer_path(start, end)
        min_path = []
        dirc = []
        train = []
        nonstop_stations = []
        total_time = 0
        i = 0
        for j in range(len(min_transfer_path) - 1):
            path = nonstop(self.data, min_transfer_path[i], min_transfer_path[j + 1])[1]
            if not path:
                time, path, new_dirc, new_train = nonstop(
                    self.data, min_transfer_path[i], min_transfer_path[j]
                )
                total_time += time
                min_path += path[:-1]
                nonstop_stations.append(min_transfer_path[j])
                dirc.append(new_dirc)
                train.append(new_train)
                i = j
        time, path, new_dirc, new_train = nonstop(self.data, min_transfer_path[i], end)
        total_time += time
        min_path += path
        dirc.append(new_dirc)
        train.append(new_train)
        time -= 10 * len(nonstop_stations)
        return total_time, min_path, nonstop_stations, dirc, train

# def main():
#     data = load_data("MergedTrainInfo_new.json")

#     railway = Railway(data)

#     """result = nonstop(data, "鲅鱼圈", "天津西")
#     if result == None:
#         print("没有找到路径")
#     else:
#         time, path, dire, train = result
#         print("最短时间: ", time)
#         print("路径:", end=" ")
#         print("->".join(path))
#         print("方向: ", dire)
#         print("列车: ", train)"""

#     time, path, station, dire, train = railway.shortest_time_path_with_nonstop_check(
#         "菏泽", "沈阳"
#     )  # 测试函数
#     if path:
#         print("最短时间: ", time)
#         print("路径: ")
#         print("->".join(path))
#         print("换乘站: ", station)
#         print("方向: ", dire)
#         print("列车: ", train)
#     else:
#         print("没有找到路径")

# if __name__ == "__main__":
#     main()
